import pandas as pd
import xgboost as xgb
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
from sklearn.model_selection import GridSearchCV
import os

# 1. 读取训练集和测试集数据
# 请将下面的路径替换为你的实际文件路径
train_data_path = "E:\\heart.csv"  # 训练集路径
test_data_path = "D:\\heart_cleaned.csv"  # 测试集路径

# 确认文件是否存在
if not os.path.exists(train_data_path) or not os.path.exists(test_data_path):
    raise FileNotFoundError("请确认训练集或测试集文件路径是否正确！")

train_data = pd.read_csv(train_data_path)  # 读取训练集
test_data = pd.read_csv(test_data_path)  # 读取测试集

# 2. 分离训练集的特征和目标列
X_train = train_data.drop(columns=['target'])  # 训练特征
y_train = train_data['target']  # 训练目标

# 确保测试集的特征和训练集的列名一致
X_test = test_data[X_train.columns]

# 3. XGBoost 参数优化
param_grid = {
    'n_estimators': [100, 200],      # 树的数量
    'max_depth': [3, 4, 5],          # 树的最大深度
    'learning_rate': [0.01, 0.1],    # 学习率
    'subsample': [0.8, 1.0],         # 每次迭代抽取的样本比例
    'colsample_bytree': [0.8, 1.0],  # 每次迭代抽取的特征比例
}

xgb_clf = xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss')

# 使用 GridSearchCV 进行超参数搜索，禁用并行计算
grid_search = GridSearchCV(estimator=xgb_clf, param_grid=param_grid, cv=5, scoring='accuracy', verbose=2, n_jobs=1)
grid_search.fit(X_train, y_train)

# 输出最佳参数
print("最佳参数:", grid_search.best_params_)

# 4. 使用最佳参数重新训练模型
best_model = grid_search.best_estimator_  # 获取最佳模型

# 5. 模型预测
y_pred = best_model.predict(X_test)

# 6. 保存预测结果到CSV文件
output = X_test.copy()  # 拷贝测试特征数据
output['predicted_target'] = y_pred  # 添加预测结果列

# 如果测试集包含真实的目标列，保留真实值以便对比
if 'target' in test_data.columns:
    output['actual_target'] = test_data['target']

# 请将下面的路径替换为你的实际文件路径
output.to_csv('D:\\heart_optimized_predictions.csv', index=False, encoding='utf-8')  # 保存结果
print("预测结果已保存为 'heart_optimized_predictions.csv'")

# 7. 模型性能评估（基于训练集）
y_train_pred = best_model.predict(X_train)

print("\n模型在训练集上的表现：")
print("准确率：", accuracy_score(y_train, y_train_pred))  # 准确率
print("\n混淆矩阵：")
print(confusion_matrix(y_train, y_train_pred))  # 混淆矩阵
print("\n分类报告：")
print(classification_report(y_train, y_train_pred))  # 分类报告